Implement new experimental lookup-based matrix multiplication method(TMAC) #26695
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Liqun Fu <[email protected]>
Signed-off-by: Liqun Fu <[email protected]>
Signed-off-by: Liqun Fu <[email protected]>
Signed-off-by: Liqun Fu <[email protected]>
Signed-off-by: Liqun Fu <[email protected]>
Signed-off-by: Liqun Fu <[email protected]>
Signed-off-by: Liqun Fu <[email protected]>
…as kernel not implemented for fp32. Also, I need to write the packing logic for the scales as well.
…ssert issue with the data shuffling in prepack
| auto scale_ptr = scales ? scales->DataRaw() : nullptr; | ||
| packed_b_ = IAllocator::MakeUniquePtr<void>(alloc, packed_b_size_, true); | ||
| MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type_, qptr, packed_b_.get(), scale_ptr, | ||
| has_zp_input_, nullptr, threadpool_ptr); |
Member
There was a problem hiding this comment.
IIUC - The usage of threadpool in the existing non-LUT path seems like a new addition - is that intentaional (and come with apprioriate tests) ?
Contributor
Author
There was a problem hiding this comment.
Initially, I thought tests in test_sqnbitgemm.cpp should suffice since they already test it with thread pool. I applied changes to only use thread pool for LUT path now.
Once we add tests, I think it might be beneficial to use thread pool for pre packing for other paths
Contributor
There was a problem hiding this comment.
closing comment for now to merge as discussed offline
| } | ||
|
|
||
| // Conditional pragma unroll for compiler compatibility | ||
| #if defined(__INTEL_COMPILER) || defined(__clang__) |
Member
There was a problem hiding this comment.
Why is this complier dependent ? Is this implementation from the T-MAC library as is ?
| // Each iteration processes one row of the activation matrix | ||
| // TODO(vraspar): Ideally we have to do block parallelism here | ||
|
|
||
| MlasTrySimpleParallel( |
Member
There was a problem hiding this comment.
If M == 1, can we parallelize on N ?
hariharans29
approved these changes
Jan 15, 2026
| } | ||
|
|
||
| size_t n_div = 0; | ||
| switch (BlkBitWidth) { |
Member
There was a problem hiding this comment.
Why have this switch if BlkBitWidth is guaranteed to be 2 at this stage ?
Contributor
Author
There was a problem hiding this comment.
I decided to have it generalized for when we add int4 kernels
| * @brief Parameters for TMAC kernel | ||
| */ | ||
| struct MlasTMACKernelParams { | ||
| size_t g; |
Member
There was a problem hiding this comment.
A brief comment describing what each config is and what it is used for will help
| ) | ||
| { | ||
| const MlasTMACKernelParams& tmac_params = MlasGetLutGemmKernelParams(N, K, BlkBitWidth, BlkLen, HasZeroPoint); | ||
| const size_t PackedQuantBDataSize = (N * BlkBitWidth) * (K / tmac_params.g / tmac_params.ngroups_per_elem); |
Member
There was a problem hiding this comment.
Is there an alignment requirement for the packed weights ?
| assert(bm % mgroup == 0); | ||
| assert(bm % bits == 0); | ||
|
|
||
| std::unique_ptr<uint8_t[]> buf(new uint8_t[N * bits * (K / g)]); |
Member
There was a problem hiding this comment.
For the purpose of what is being done here, a standard RAII containers like vector would do, do we really need a unique_ptr here ?
alex-spacemit
pushed a commit
to spacemit-com/onnxruntime
that referenced
this pull request
Jan 20, 2026
…TMAC) (microsoft#26695) ### Description This PR introduces a new experimental lookup-table(LUT) based matrix multiplication method for 2-bit MatMulNBits on x64 AVX2 inspired from [T-MAC paper](https://arxiv.org/abs/2407.00088) and [T-MAC repository](https://github.com/microsoft/T-MAC) to speed up low bit LLM inference. Unlike the existing quant-dequant methods, the LUT-based method directly supports mixed-precision-GEMM without dequantization. It uses bit-wise table lookup to eliminate multiplications and reduce additions required in matrix multiplication. <img width="1910" height="759" alt="image" src="https://github.com/user-attachments/assets/3e3f2ced-eba4-4d4e-a63c-fec479943202" /> This PR: - Add` mlas.use_lut_gemm` session option allowing use of LUT GEMM inside matmulnbits when it is available (2-bit, BlkLen multiple of 32, K multiple of 32, N multiple of 128, AVX2 present). - Introduces LUT packing + kernel config cache (packs bitplanes, scales, ZP) and the main `MlasLUTGemm` entry that generates per-row LUTs and calls the AVX2 kernel. - Implements AVX2 LUT generation `GenerateLUT_avx2` and GEMM compute `TMACComputeGemm_avx2` and wires dispatch in MLAS platform init. - Updates MatMulNBits PrePack/Compute to use LUT packing/compute when opted-in; keeps existing quant-dequant path as fallback. - Extends Python quant bindings with 2-bit QDQ helper for parity with the new path. - Adds MLAS unit tests covering LUT GEMM across symmetric/asymmetric quant and multiple shapes/block sizes. ### Main components: - `MlasInitLUTGemmKernelConfig`: Config for LUT kernels - `MlasLUTGemmPackQuantBData`: Pre Packing of quantized weight - `MlasLUTPackScalesAndZeroPoints`: Pre Packing of qunatized scales and zero points - `MlasLUTGemm`: Main Entry point - `GenerateLUT_avx2`: LUT construction from activations - `TMACComputeGemm_avx2`: AVX2 LUT GEMM kernel - Session option: mlas.use_lut_gemm ### How to test - MLAS LUT GEMM unit tests: see `test_sqlutgemm.cpp` - Run MatMulNBits models with session option `mlas.use_lut_gemm=1` on AVX2 machines; expect fallback to existing path if availability checks fail. ### Perf Focus of this PR is functional + kernel bring-up; perf to be reported separately once broader profiling is done. ### Future Work - Support MLFloat16 (FP16 scales and zero points) - Add neon kernel for ARM. - Add kernels for 4 bit weights and bitnet kernels - Broader batch (N>1) support and additional shape coverage. --------- Signed-off-by: Liqun Fu <[email protected]> Co-authored-by: Liqun Fu <[email protected]> Co-authored-by: carzh <[email protected]> Co-authored-by: Hector Li <[email protected]> Co-authored-by: carzh <[email protected]> Co-authored-by: Vrajang Parikh <[email protected]>
tianleiwu
pushed a commit
that referenced
this pull request
Jan 21, 2026
…TMAC) (#26695) ### Description This PR introduces a new experimental lookup-table(LUT) based matrix multiplication method for 2-bit MatMulNBits on x64 AVX2 inspired from [T-MAC paper](https://arxiv.org/abs/2407.00088) and [T-MAC repository](https://github.com/microsoft/T-MAC) to speed up low bit LLM inference. Unlike the existing quant-dequant methods, the LUT-based method directly supports mixed-precision-GEMM without dequantization. It uses bit-wise table lookup to eliminate multiplications and reduce additions required in matrix multiplication. <img width="1910" height="759" alt="image" src="https://github.com/user-attachments/assets/3e3f2ced-eba4-4d4e-a63c-fec479943202" /> This PR: - Add` mlas.use_lut_gemm` session option allowing use of LUT GEMM inside matmulnbits when it is available (2-bit, BlkLen multiple of 32, K multiple of 32, N multiple of 128, AVX2 present). - Introduces LUT packing + kernel config cache (packs bitplanes, scales, ZP) and the main `MlasLUTGemm` entry that generates per-row LUTs and calls the AVX2 kernel. - Implements AVX2 LUT generation `GenerateLUT_avx2` and GEMM compute `TMACComputeGemm_avx2` and wires dispatch in MLAS platform init. - Updates MatMulNBits PrePack/Compute to use LUT packing/compute when opted-in; keeps existing quant-dequant path as fallback. - Extends Python quant bindings with 2-bit QDQ helper for parity with the new path. - Adds MLAS unit tests covering LUT GEMM across symmetric/asymmetric quant and multiple shapes/block sizes. ### Main components: - `MlasInitLUTGemmKernelConfig`: Config for LUT kernels - `MlasLUTGemmPackQuantBData`: Pre Packing of quantized weight - `MlasLUTPackScalesAndZeroPoints`: Pre Packing of qunatized scales and zero points - `MlasLUTGemm`: Main Entry point - `GenerateLUT_avx2`: LUT construction from activations - `TMACComputeGemm_avx2`: AVX2 LUT GEMM kernel - Session option: mlas.use_lut_gemm ### How to test - MLAS LUT GEMM unit tests: see `test_sqlutgemm.cpp` - Run MatMulNBits models with session option `mlas.use_lut_gemm=1` on AVX2 machines; expect fallback to existing path if availability checks fail. ### Perf Focus of this PR is functional + kernel bring-up; perf to be reported separately once broader profiling is done. ### Future Work - Support MLFloat16 (FP16 scales and zero points) - Add neon kernel for ARM. - Add kernels for 4 bit weights and bitnet kernels - Broader batch (N>1) support and additional shape coverage. --------- Signed-off-by: Liqun Fu <[email protected]> Co-authored-by: Liqun Fu <[email protected]> Co-authored-by: carzh <[email protected]> Co-authored-by: Hector Li <[email protected]> Co-authored-by: carzh <[email protected]> Co-authored-by: Vrajang Parikh <[email protected]> (cherry picked from commit 8e050d1)
This was referenced Jan 21, 2026
tianleiwu
added a commit
that referenced
this pull request
Jan 23, 2026
### Description This PR cherry-picks the following changes for the 1.24.0 release. ### Cherry-picked Commits | Commit | Commit Title | Author | |---|---|---| | 744e7fe | Add type definitions, registration, utilities for INT2/UINT2 support (#26824) | vraspar | | 530a1fb | [QNN EP] Add BFloat16 dtype support in QNN EP (#26987) | tirupath-qti | | 8e050d1 | Implement new experimental lookup-based matrix multiplication method(TMAC) (#26695) | vraspar | | 2d2ba6b | [MLAS/CPU EP] Improve performance of Silu activation path within the QuickGelu CPU kernel (#26753) | Hariharan Seshadri | | 1c02b79 | [QNN EP] Add support for handling 0-dimension for Concat Op (#27000) | Ashwath Shankarnarayan | | cc2b01b | Fix ClipQuantFusion crash when Clip has multiple input edges (#27016) | Edward Chen | | bbd3850 | [QNN EP] Support quantized BatchNorm with per-channel DQ params on QNN HTP (#26959) | qti-yuduo | | d8f0318 | Add API to get ep graph partitioning info (#26781) | Adrian Lizarraga | | b912b18 | [OVEP] OpenVINO EP Features and bug-fixes for ORT-1.24 - Follow up (#27007) | Preetha Veeramalai | | ba11af4 | [QNN-EP] Add MatMulNBits translation for GPU (#26340) | quic-tirupath | | c03c419 | [MLAS/NEON] Add dedicated kernel for depthwise convolution for ARM64 using NEON intrinsics (#26688) | Hariharan Seshadri | | e7dfd69 | [QNN-EP] Support alternate Layernorm fusion pattern in QNN preprocess (#26060) | qti-mattsinc | | 4013dc1 | Implement multithreading in qgemm_kleidi (#26301) | Melike Kaptan | | 9f06181 | [CXX] Enable users to specify custom OrtSyncStream via RunOptions (#26988) | Dmitri Smirnov | | cfccd64 | Added support for QMX kernels in MLAS (#26849) | qti-vaiskv | | 29d9b2f | Tweak external resource importer handle structs (#27040) | Scott McKay | | 9d108d0 | [QNN EP] Add QuickGELU operator support for QNN provider (#27034) | tirupath-qti | | b35688f | Add INT2 and UINT2 support for QDQ, transpose and cast ops (#27022) | vraspar | | 6d34aba | Introducing BF16 Pointwise NCHWc Convolution for Arm64 (#26838) | Rohanjames1997 | | 36017ad | [EP ABI] Add CreateCustomOpDomains() API for plugin EP to register custom ops (#27050) | Chi Lo | | 50a03e4 | Add a new pipeline for CUDA 13 nuget builds (#27023) | eserscor | | a0d4439 | [EP ABI] Update Graph_GetGraphView() implementation (#26711) | Chi Lo | | 34bb209 | [webgpu] Fix a bug for im2col (#27069) | Wenqin Yang | | 46e8d45 | [QNN EP] Add FusedMatMul operator support (#27044) | tirupath-qti | | 5e7e7a3 | Disable Float32_2Bits_Asymmetric_256x256 test (#27046) | vraspar | | 39f966e | Fix Doxygen documentation build error in onnxruntime_c_api.h (#27083) | Nick Eubanks | | 8a7a797 | Print tensor for new packed type of 2 bits (#27064) | Tianlei Wu | | 01f40e6 | Fix GPU JAR testing on Linux (#27011) | eserscor | | b6ed7f3 | Fix warning around ununsed code in QNN Android Emulator builds by clang (#27026) | Hariharan Seshadri | | d7daa45 | Raise the timeout for the ios simulator job (#27045) | Hariharan Seshadri | | 7e1d818 | upgrade emsdk to 4.0.23 (#27029) | Yulong Wang | | 347b990 | Fix failing mainline build on Arm64 linux (#27101) | Rohanjames1997 | | f481b17 | Add dedicated API to support extracting compatibility string from model metadata (#27015) | adrastogi | --------- Signed-off-by: Liqun Fu <[email protected]> Signed-off-by: bfilipek <[email protected]> Signed-off-by: dependabot[bot] <[email protected]> Signed-off-by: Jonathan Clohessy <[email protected]> Signed-off-by: Christian Bourjau <[email protected]> Signed-off-by: melkap01 <[email protected]> Co-authored-by: vraspar <[email protected]> Co-authored-by: tirupath-qti <[email protected]> Co-authored-by: Ashwath Shankarnarayan <[email protected]> Co-authored-by: Liqun Fu <[email protected]> Co-authored-by: carzh <[email protected]> Co-authored-by: Hector Li <[email protected]> Co-authored-by: carzh <[email protected]> Co-authored-by: Vrajang Parikh <[email protected]> Co-authored-by: Hariharan Seshadri <[email protected]> Co-authored-by: Copilot <[email protected]> Co-authored-by: Edward Chen <[email protected]> Co-authored-by: Yuduo Wu <[email protected]> Co-authored-by: Adrian Lizarraga <[email protected]> Co-authored-by: Preetha Veeramalai <[email protected]> Co-authored-by: jatinwadhwa921 <[email protected]> Co-authored-by: jatinwadhwa921 <[email protected]> Co-authored-by: saurabh <[email protected]> Co-authored-by: Ankit Maheshkar <[email protected]> Co-authored-by: sfatimar <[email protected]> Co-authored-by: Javier Martinez <[email protected]> Co-authored-by: Bartlomiej Filipek <[email protected]> Co-authored-by: bopeng1234 <[email protected]> Co-authored-by: Eric Crawford <[email protected]> Co-authored-by: MayureshV1 <[email protected]> Co-authored-by: TejalKhade28 <[email protected]> Co-authored-by: Vishnudas Thaniel S <[email protected]> Co-authored-by: Yaru Du <[email protected]> Co-authored-by: Ryan Metcalfe <[email protected]> Co-authored-by: Dvoretckii, Mikhail <[email protected]> Co-authored-by: Pallavi Gupta <[email protected]> Co-authored-by: Jianhui Dai <[email protected]> Co-authored-by: Jiajia Qin <[email protected]> Co-authored-by: Changming Sun <[email protected]> Co-authored-by: Fei Chen <[email protected]> Co-authored-by: Yulong Wang <[email protected]> Co-authored-by: Akupadhye <[email protected]> Co-authored-by: Wang Ning <[email protected]> Co-authored-by: Maximilian Müller <[email protected]> Co-authored-by: Chi Lo <[email protected]> Co-authored-by: George Wu <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Wanming Lin <[email protected]> Co-authored-by: quic-calvnguy <[email protected]> Co-authored-by: Jie Chen <[email protected]> Co-authored-by: xhcao <[email protected]> Co-authored-by: Wei-Sheng Chin <[email protected]> Co-authored-by: quic-hungjuiw <[email protected]> Co-authored-by: Ian Hunter <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: kunal-vaishnavi <[email protected]> Co-authored-by: Jeff Kilpatrick <[email protected]> Co-authored-by: Jeff Kilpatrick <[email protected]> Co-authored-by: Scott McKay <[email protected]> Co-authored-by: Nenad Banfic <[email protected]> Co-authored-by: derdeljan-msft <[email protected]> Co-authored-by: n1harika <[email protected]> Co-authored-by: Ryan Metcalfe <[email protected]> Co-authored-by: Jaswanth Gannamaneni <[email protected]> Co-authored-by: Klimenko, Mikhail <[email protected]> Co-authored-by: liang <[email protected]> Co-authored-by: Garth Long <[email protected]> Co-authored-by: Jonathan Clohessy <[email protected]> Co-authored-by: Akshay Sonawane <[email protected]> Co-authored-by: Christopher Warrington <[email protected]> Co-authored-by: Ishwar Raut <[email protected]> Co-authored-by: Gaurav Garg <[email protected]> Co-authored-by: Xinpeng Dou <[email protected]> Co-authored-by: adrastogi <[email protected]> Co-authored-by: Aditya Rastogi <[email protected]> Co-authored-by: qti-hungjuiw <[email protected]> Co-authored-by: Pradeep Sakhamoori <[email protected]> Co-authored-by: Adam Pocock <[email protected]> Co-authored-by: mingyue <[email protected]> Co-authored-by: Susanta Bhattacharjee <[email protected]> Co-authored-by: Jozef Wludzik <[email protected]> Co-authored-by: Rajeev Sekar <[email protected]> Co-authored-by: Mayuresh M Varerkar <[email protected]> Co-authored-by: Copilot <[email protected]> Co-authored-by: Wenqin Yang <[email protected]> Co-authored-by: xieofxie <[email protected]> Co-authored-by: hualxie <[email protected]> Co-authored-by: Joshua Lochner <[email protected]> Co-authored-by: Christian Bourjau <[email protected]> Co-authored-by: Xiaofei Han <[email protected]> Co-authored-by: Dmitri Smirnov <[email protected]> Co-authored-by: chunghow-qti <[email protected]> Co-authored-by: Guenther Schmuelling <[email protected]> Co-authored-by: Jiawei Shao <[email protected]> Co-authored-by: czekun <[email protected]> Co-authored-by: Jaskaran Singh Nagi <[email protected]> Co-authored-by: quic-tirupath <[email protected]> Co-authored-by: qti-mattsinc <[email protected]> Co-authored-by: Melike Kaptan <[email protected]> Co-authored-by: Damien Dooley <[email protected]> Co-authored-by: qti-vaiskv <[email protected]> Co-authored-by: Rohanjames1997 <[email protected]> Co-authored-by: eserscor <[email protected]> Co-authored-by: eserscor <[email protected]> Co-authored-by: Nick Eubanks <[email protected]> Co-authored-by: adrastogi <[email protected]> Co-authored-by: Rohanjames1997 <[email protected]>
tianleiwu
added
cherry-picked
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR introduces a new experimental lookup-table(LUT) based matrix multiplication method for 2-bit MatMulNBits on x64 AVX2 inspired from T-MAC paper and T-MAC repository to speed up low bit LLM inference.
Unlike the existing quant-dequant methods, the LUT-based method directly supports mixed-precision-GEMM without dequantization. It uses bit-wise table lookup to eliminate multiplications and reduce additions required in matrix multiplication.
This PR:
mlas.use_lut_gemmsession option allowing use of LUT GEMM inside matmulnbits when it is available (2-bit, BlkLen multiple of 32, K multiple of 32, N multiple of 128, AVX2 present).MlasLUTGemmentry that generates per-row LUTs and calls the AVX2 kernel.GenerateLUT_avx2and GEMM computeTMACComputeGemm_avx2and wires dispatch in MLAS platform init.Main components:
MlasInitLUTGemmKernelConfig: Config for LUT kernelsMlasLUTGemmPackQuantBData: Pre Packing of quantized weightMlasLUTPackScalesAndZeroPoints: Pre Packing of qunatized scales and zero pointsMlasLUTGemm: Main Entry pointGenerateLUT_avx2: LUT construction from activationsTMACComputeGemm_avx2: AVX2 LUT GEMM kernelSession option: mlas.use_lut_gemm
How to test
test_sqlutgemm.cppmlas.use_lut_gemm=1on AVX2 machines; expect fallback to existing path if availability checks fail.Perf
Focus of this PR is functional + kernel bring-up; perf to be reported separately once broader profiling is done.
Future Work